package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"github.com/cockroachdb/pebble"
	"github.com/gin-gonic/gin"
	"github.com/klauspost/compress/zstd"
	"golang.org/x/sync/errgroup"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"
	"unsafe"
)

var (
	chunkSize = 1 << 20
)

// Metadata 定义文件的元数据结构
type Metadata struct {
	Name   string    `json:"name"`
	Size   int64     `json:"size"`
	Type   string    `json:"type"`
	Time   time.Time `json:"time"`
	Chunks []string  `json:"chunks"`
}

// splitFile 将文件分割成1kb的分片
func splitFile(filePath string) ([]string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var chunks []string
	var g errgroup.Group
	for {
		buffer := make([]byte, chunkSize)
		read, err := file.Read(buffer)
		if err != nil && err != io.EOF {
			return nil, err
		}

		if read > 0 {
			chunk := buffer[:read]
			hashCode := Sum15BitSHA256(chunk)
			chunks = append(chunks, hashCode)
			g.Go(func() error {
				compressedChunk, err := compressChunk(chunk)
				if err != nil {
					return err
				}
				err = saveFile(hashCode, compressedChunk)
				if err != nil {
					return err
				}
				return nil
			})
		}

		if err == io.EOF {
			break
		}
	}
	err = g.Wait()
	if err != nil {
		return nil, err
	}
	return chunks, nil
}

func getFileType(path string) (string, error) {
	file, err := os.Open(path)
	if err != nil {
		return "", err
	}
	buf := make([]byte, 512)
	_, err = file.Read(buf)
	if err != nil {
		return "", err
	}
	return http.DetectContentType(buf), nil
}

// compressChunk 使用zstd压缩分片
func compressChunk(chunk []byte) ([]byte, error) {
	var b bytes.Buffer
	w, err := zstd.NewWriter(&b)
	if err != nil {
		return nil, err
	}

	_, err = w.Write(chunk)
	if err != nil {
		return nil, err
	}

	err = w.Close()
	if err != nil {
		return nil, err
	}

	return b.Bytes(), nil
}

func decompressChunk(chunkPath string) ([]byte, error) {
	chunkFile, err := os.ReadFile(chunkPath)
	if err != nil {
		return nil, err
	}
	d, err := zstd.NewReader(nil)
	if err != nil {
		return nil, err
	}
	defer d.Close()
	res, err := d.DecodeAll(chunkFile, nil)
	if err != nil {
		return nil, err
	}
	return res, nil
}

// saveMetadata 将元数据保存到pebbleDB
func saveMetadata(db *pebble.DB, metadata *Metadata) error {
	metaBytes, err := json.Marshal(metadata)
	if err != nil {
		return err
	}

	return db.Set(StringToBytes(metadata.Name), metaBytes, &pebble.WriteOptions{})
}

// loadMetadata 从pebbleDB加载元数据
func loadMetadata(db *pebble.DB, fileName string) (*Metadata, error) {
	metaBytes, closer, err := db.Get(StringToBytes(fileName))
	if err != nil {
		return nil, err
	}
	defer closer.Close()

	var metadata Metadata
	err = json.Unmarshal(metaBytes, &metadata)
	if err != nil {
		return nil, err
	}

	return &metadata, nil
}

func StringToBytes(s string) []byte {
	return unsafe.Slice(unsafe.StringData(s), len(s))
}

func BytesToString(b []byte) string {
	return unsafe.String(&b[0], len(b))
}

func getSHA256(data []byte) []byte {
	// 创建 SHA256 哈希对象
	hash := sha256.New()

	// 将数据写入哈希对象
	hash.Write(data)

	// 计算哈希值
	sha256Hash := hash.Sum(nil)

	return sha256Hash
}

func getSHA256String(hash []byte) string {
	hashString := hex.EncodeToString(hash)
	return strings.ToUpper(hashString[:15])
}

func Sum15BitSHA256(data []byte) string {
	return getSHA256String(getSHA256(data))
}

func saveFile(hashCode string, data []byte) error {
	// 构建文件路径
	filePath := getPathFromHash(hashCode)

	// 创建目录
	err := os.MkdirAll(filepath.Dir(filePath), 0755)
	if err != nil {
		return err
	}

	// 写入文件
	err = os.WriteFile(filePath, data, 0644)
	if err != nil {
		return err
	}

	return nil
}

func getPathFromHash(hash string) string {
	return filepath.Join(".", "data", string(hash[0]), hash[1:3], hash[3:])
}

type App struct {
	db *pebble.DB
}

func main() {
	// 初始化pebbleDB
	db, err := pebble.Open("filestore.db", &pebble.Options{})
	if err != nil {
		panic(err)
	}
	defer db.Close()
	app := App{
		db,
	}
	r := gin.Default()
	//r.MaxMultipartMemory = 8 << 20 // 8 MiB
	r.POST("/upload", app.uploadFile)
	r.GET("/download/:filename", app.downloadFile)

	if err = r.Run(":8080"); err != nil {
		log.Println(err)
	}
}

func (a *App) uploadFile(c *gin.Context) {
	file, err := c.FormFile("file")
	if err != nil {
		c.String(http.StatusBadRequest, fmt.Sprintf("上传文件失败: %s", err.Error()))
		return
	}

	var metadata Metadata
	metadata.Name = file.Filename
	metadata.Size = file.Size // 假设文件只有一个分片
	//metadata.Type = fileType        // 根据文件扩展名确定文件类型
	metadata.Time = time.Now()
	filename := file.Filename
	open, err := file.Open()
	if err != nil {
		c.String(http.StatusBadRequest, fmt.Sprintf("上传文件失败: %s", err.Error()))
		return
	}

	var flg = false
	var g errgroup.Group
	for {
		buffer := make([]byte, chunkSize)
		read, err := open.Read(buffer)
		if err != nil && err != io.EOF {
			c.String(http.StatusBadRequest, fmt.Sprintf("上传文件失败: %s", err.Error()))
			return
		}

		if read > 0 {
			chunk := buffer[:read]
			if !flg {
				flg = true
				metadata.Type = http.DetectContentType(chunk)
			}
			hashCode := Sum15BitSHA256(chunk)
			metadata.Chunks = append(metadata.Chunks, hashCode)
			g.Go(func() error {
				compressedChunk, err := compressChunk(chunk)
				if err != nil {
					return err
				}
				err = saveFile(hashCode, compressedChunk)
				if err != nil {
					return err
				}
				return nil
			})
		}

		if err == io.EOF {
			break
		}
	}
	err = g.Wait()
	if err != nil {
		c.String(http.StatusBadRequest, "分片存储失败：%s\n", err)
		return
	}
	err = saveMetadata(a.db, &metadata)
	if err != nil {
		c.String(http.StatusBadRequest, "储存元数据失败：%s\n", err)
		return
	}

	c.String(http.StatusOK, fmt.Sprintf("文件 %s 上传成功", filename))
}

func (a *App) downloadFile(c *gin.Context) {
	filename := c.Param("filename")
	loadedMetadata, err := loadMetadata(a.db, filename)
	if err != nil {
		c.String(http.StatusBadRequest, fmt.Sprintf("获取元数据失败: %s", err.Error()))
		return
	}
	var paths []string
	for _, chunk := range loadedMetadata.Chunks {
		chunkFilePath := getPathFromHash(chunk)
		paths = append(paths, chunkFilePath)
	}
	extraHeaders := map[string]string{
		"Last-Modified": loadedMetadata.Time.UTC().Format(http.TimeFormat),
		"Content-Disposition": fmt.Sprintf(`attachment; filename="%s"; filename*=UTF-8''%s`,
			url.QueryEscape(loadedMetadata.Name), url.QueryEscape(loadedMetadata.Name),
		),
	}
	data, err := getMultiDecompressedReader(paths)
	if err != nil {
		c.String(http.StatusBadRequest, fmt.Sprintf("创建文件流失败: %s", err.Error()))
		return
	}
	c.DataFromReader(http.StatusOK, loadedMetadata.Size, "application/octet‑stream", data, extraHeaders)
}

func getMultiDecompressedReader(filePaths []string) (io.Reader, error) {
	readers := make([]io.Reader, len(filePaths))
	for i, filePath := range filePaths {
		file, err := os.Open(filePath)
		if err != nil {
			return nil, err
		}

		zr, err := zstd.NewReader(file)
		if err != nil {
			file.Close()
			return nil, err
		}

		readers[i] = zr
	}

	return io.MultiReader(readers...), nil
}
